package com.zhilei.deepseekdoctor.controller;

import com.zhilei.deepseekdoctor.pojo.ChatEntity;
import com.zhilei.deepseekdoctor.pojo.ChatRecord;
import com.zhilei.deepseekdoctor.service.OllamaService;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.stream.Collectors;

/**
 * 连接ollama的控制器
 *
 * @author zhilei
 * @version 1.0
 */
@RestController
@RequestMapping("/api/ollama")
public class OllamaController {

    @Resource
    private OllamaService ollamaService;

    @Resource
    OllamaChatClient ollamaChatClient;

    /**
     * 调用ollama的接口
     * 同步调用
     * 会阻塞当前进程，直到数据返回
     * 传入的是一个提示词
     *
     * @param msg
     * @return
     */
    @GetMapping("ai/chat")
    public String aiOllamaChat(@RequestParam String msg) {
        return ollamaChatClient.call(msg);
    }


    /**
     * 调用ollama的接口
     * 流式调用
     * 不会阻塞当前进程，会立即返回一个任务id
     * 传入的是一个提示词对象
     *
     * @param msg
     * @return
     */
    @GetMapping("ai/chat/stream")
    public Flux<ChatResponse> aiOllamaChatStream(@RequestParam String msg) {
        UserMessage userMessage = new UserMessage(msg);
        Prompt prompt = new Prompt(userMessage);
        return ollamaChatClient.stream(prompt);
    }

    /**
     * 调用ollama的接口
     * 流式调用
     * 不会阻塞当前进程，会立即返回一个任务id
     * 传入的是一个提示词对象
     *
     * @param msg
     * @return
     */
    @GetMapping("ai/chat/streamCount")
    public List<String> aiOllamaChatstreamCount(@RequestParam String msg) {
        UserMessage userMessage = new UserMessage(msg);
        Prompt prompt = new Prompt(userMessage);
        Flux<ChatResponse> streamResponse = ollamaChatClient.stream(prompt);

        List<String> collect = streamResponse.toStream().map(s -> s.getResult().getOutput().getContent())
                .collect((Collectors.toList()));

        return collect;
    }

    /**
     * 调用ollama的接口
     * 同步调用
     * 会阻塞当前进程，直到数据返回
     * 传入的是一个提示词
     *
     * @param msg
     * @return
     */
    @GetMapping("ai/v2/chat")
    public String aiV2OllamaChat(@RequestParam String msg) {
        return ollamaService.aiOllamaChat(msg);
    }


    /**
     * 调用ollama的接口
     * 流式调用
     * 不会阻塞当前进程，会立即返回一个任务id
     * 传入的是一个提示词对象
     *
     * @param msg
     * @return
     */
    @GetMapping("ai/v2/chat/stream")
    public Flux<ChatResponse> aiV2OllamaChatStream(@RequestParam String msg) {
        return ollamaService.aiOllamaChatStream(msg);
    }

    /**
     * 调用ollama的接口
     * 流式调用
     * 不会阻塞当前进程，会立即返回一个任务id
     * 传入的是一个提示词对象
     *
     * @param msg
     * @return
     */
    @GetMapping("ai/v2/chat/streamCount")
    public List<String> aiV2OllamaChatstreamCount(@RequestParam String msg) {
        return ollamaService.aiOllamaChatstreamCount(msg);
    }

    /**
     * 调用ollama的接口
     * 流式调用
     * 不会阻塞当前进程，会立即返回一个任务id
     * 传入的是一个提示词对象
     *
     * @param chat
     * @return
     */
    @PostMapping("ai/v3/chat/aiV3OllamaChatstreamCount")
    public void aiV3OllamaChatstreamCount(@RequestBody ChatEntity chat) {
        ollamaService.aiV3OllamaChatstreamCount(chat);
    }

    /**
     * 保存当前的沟通的对象
     *
     * @param who
     * @return
     */
    @GetMapping("ai/v3/ollama/getRecords")
    public List<ChatRecord> getRecords(@RequestParam String who) {
        return ollamaService.getRecords(who);
    }



}



































